Skip to content

fix(serve): make torch an optional dependency - #6166

Open
goelakash wants to merge 1 commit into
aws:masterfrom
goelakash:fix-optional-torch-serve
Open

fix(serve): make torch an optional dependency#6166
goelakash wants to merge 1 commit into
aws:masterfrom
goelakash:fix-optional-torch-serve

Conversation

@goelakash

Copy link
Copy Markdown
Contributor

Fixes #5531

sagemaker-serve lists torch>=2.0.0 as a required dependency, so any install of sagemaker, sagemaker-serve, or sagemaker-mlops pulls torch (and on Linux, the CUDA/cuDNN/NCCL stack) even for API-only use. sagemaker-core already treats torch as an extra.

The blocker was DEFAULT_SERIALIZERS_BY_FRAMEWORK in serve/constants.py, which instantiated TorchTensorSerializer() at module scope. That runs from torch import Tensor on any import sagemaker.serve. Every other torch reference in the package is already lazy.

Changes

  • Store serializer/deserializer classes in DEFAULT_SERIALIZERS_BY_FRAMEWORK; instantiate at lookup in _fetch_serializer_and_deserializer_for_framework.
  • Move torch>=2.0.0 to a torch extra, matching sagemaker-core.
  • Update the two docstring examples and the affected unit tests.
  • Add tests/unit/test_optional_torch_dependency.py, mirroring the existing sagemaker-core subprocess pattern.

Installing with the torch extra is unchanged. TorchTensorSerializer() still raises the same ImportError if torch is missing when actually used.

Testing

  • tests/unit in a venv with and without torch: no new failures vs. master (53 pre-existing failures in both).
  • New tests fail on master and pass with this change.
  • Verified in a container with no torch installed: from sagemaker.serve import ModelBuilder and import sagemaker.mlops both succeed.
  • Image size for all four packages on Amazon Linux 2023: 1.45 GB -> 536 MB.

Installing sagemaker-serve (or the umbrella sagemaker, which depends on
it) pulled torch, and on Linux the ~2.9 GB CUDA/cuDNN/NCCL closure, even
for API-only use. sagemaker-core already declared torch an extra;
sagemaker-serve declared it required, so serve users paid regardless.

The import-time blocker was DEFAULT_SERIALIZERS_BY_FRAMEWORK in
serve/constants.py, which instantiated TorchTensorSerializer() at module
scope, running `from torch import Tensor` on any `import sagemaker.serve`.

Changes:

- Store serializer/deserializer classes (not instances) in
  DEFAULT_SERIALIZERS_BY_FRAMEWORK; instantiate on lookup. Removes the
  import-time torch dependency.
- Move torch>=2.0.0 to a `torch` extra in sagemaker-serve, matching
  sagemaker-core. Add a `torch` extra to the umbrella sagemaker package
  (sagemaker-serve[torch]) so `pip install sagemaker[torch]` works.
- Duck-type the tensor check in TorchTensorSerializer: serialization only
  needs the object's own detach()/numpy(), so it no longer imports torch
  at all and works whenever the caller holds a tensor.
- Make the in-process model server's torch import lazy with a CPU
  fallback, so importing it no longer requires torch.
- Fix the error messages that told serve users to install the wrong
  package: TorchTensorDeserializer, the Triton translator, and the ONNX
  export path now name sagemaker-serve[torch].
- Update and add tests asserting the whole surface imports and serializes
  without torch, and that the paths that genuinely construct torch
  objects (deserializer, Triton translator) still raise a clear error
  naming the extra.

Fixes aws#5531

Breaking change: users who installed plain sagemaker-serve and relied on
a torch code path (deserializing to tensor/pt, or ONNX/Triton export of a
PyTorch model) must now install sagemaker-serve[torch]. Both paths
inherently require a torch object the caller supplies, so torch is
already present in practice; a GitHub-wide search found no external
callers of these paths that do not already import torch.

---

X-AI-Prompt: address the reviewer concern on whether making torch optional breaks existing customers; find ways to reduce the torch dependency and update pyproject and docs; consolidate all changes in the goelakash93 fork for a possible major version bump.
X-AI-Tool: Claude Code
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Make torch an optional dependency / extra

1 participant